#pragma once

#include "../RHIResource.h"
#include "../Shader/ShaderPtr.h"
#include "../Texture/RHISampler.h"
#include "../Texture/RHITexture.h"
#include "../Buffer/BufferView.h"
#include "../Raytracing/AccelerationStructure.h"
#include "../Texture/RenderTarget.h"

class DescriptorMask {
	friend class RHIDescriptor;

	uint32 setIndex;
	uint32 orderIndex;
	std::vector<VkDescriptorSet> descriptorSets;

public:

	DescriptorMask() :
		setIndex(UINT32_MAX),
		orderIndex(UINT32_MAX)
	{}
	DescriptorMask(uint32 set, uint32 order) :
		setIndex(set),
		orderIndex(order)
	{}

	const VkDescriptorSet& GetDescriptorSet() const;
};

class RHIDescriptor : public RHIResource
{
	friend class ShaderParam;

	struct DescriptorNode
	{
		VkDescriptorSetLayoutBinding descriptorSetLayoutBinding;
		VkWriteDescriptorSet writeDescriptorSet;

		std::vector<VkDescriptorImageInfo> imageInfos;
		std::vector<VkDescriptorBufferInfo> bufferInfos;
		std::vector<VkBufferView> bufferViewInfos;

		VkWriteDescriptorSetAccelerationStructureKHR tlasWrite;
		std::vector<VkAccelerationStructureKHR> tlasInfos;

		VkWriteDescriptorSetInlineUniformBlockEXT inlineWrite;

		bool isAssign;

		//##### Store rhi resource #####

		std::vector<RefCountPtr<RHITexture>> textures;
		std::vector<RefCountPtr<BufferView>> bufferViews;
		std::vector<RefCountPtr<RHIBuffer>> buffers;

		//##### Store rhi resource #####

		void ClearStore()
		{
			textures.clear();
			bufferViews.clear();
			buffers.clear();
		}


		DescriptorNode()
		{
			descriptorSetLayoutBinding = {};
			writeDescriptorSet = {};
			writeDescriptorSet.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;

			tlasWrite = {};
			tlasWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR;

			inlineWrite = {};
			inlineWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_INLINE_UNIFORM_BLOCK_EXT;

			isAssign = false;
		}
	};

	struct BuilderDetails {

		std::map<uint32, uint32> descriptorCounts;
		std::map<uint32, bool> descriptorMultiples;
		std::vector<ShaderPtr<RHIShader>> shaderStages;

		BuilderDetails()
		{
		}
	};
public:

	class Builder : public BuilderBase<BuilderDetails> {
		friend struct BuilderDetails;
	public:
		Builder() noexcept;
		Builder(Builder const& rhs) noexcept;
		Builder(Builder&& rhs) noexcept;
		~Builder() noexcept;
		Builder& operator=(Builder const& rhs) noexcept;
		Builder& operator=(Builder&& rhs) noexcept;

		Builder& SetDescriptorCount(uint32 index, uint32 count, bool isMultiple = false) noexcept;

		/************************************************************************/
		/*                           Shader                                     */
		/************************************************************************/

		Builder& PushShader(const std::string& name) noexcept;

		ShaderPtr<RHIShader> AddShader(const std::string& name) noexcept;

		template<typename ShaderType>
		Builder& PushShader() noexcept;

		template<typename ShaderType>
		ShaderPtr<ShaderType> AddShader() noexcept;

		RefCountPtr<RHIDescriptor> Build();
		RefCountPtr<RHIDescriptor> Build(RenderProcess* process, const std::string& name);

	private:

		friend class RayTracingLibrary;
	};

	RHIDescriptor(BuilderDetails& builderDetails);

	virtual void ReleaseResource() override;

	virtual ~RHIDescriptor();

	VkPipelineVertexInputStateCreateInfo& GetVertexInputStateCreateInfo();

	VkPipelineInputAssemblyStateCreateInfo& GetInputAssemblyStateCreateInfo();

	std::vector<ShaderPtr<RHIShader>>& GetShaderStages();

	VkPipelineLayout GetPipelineLayout() const;

public:

	void UpdateDescriptor(const DescriptorMask& descriptorMask, bool isMultiple = true);

	DescriptorMask AllocateDescriptor();

	ShaderParam* GetParamInfo(const std::string& name) const;

	void SetValue(const std::string& name, RefCountPtr<RHISampler> sampler);

	void SetValues(const std::string& name, const std::vector<RefCountPtr<RHISampler>>& samplers);

	void SetValue(const std::string& name, RefCountPtr<RHITexture> texture, RefCountPtr<RHISampler> sampler, VkImageLayout layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

	void SetValues(const std::string& name, const std::vector<RefCountPtr<RHITexture>>& textures, const std::vector<RefCountPtr<RHISampler>>& samplers, std::vector<VkImageLayout> layouts = std::vector<VkImageLayout>());

	void SetValue(const std::string& name, RefCountPtr<RHITexture> texture, VkImageLayout layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

	void SetValues(const std::string& name, const std::vector<RefCountPtr<RHITexture>>& textures, const std::vector<VkImageLayout> layouts = std::vector<VkImageLayout>());

	void SetValue(const std::string& name, RefCountPtr<RenderTarget> texture, VkImageLayout layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

	void SetValue(const std::string& name, RefCountPtr<BufferView> bufferView);

	void SetValues(const std::string& name, const std::vector<RefCountPtr<BufferView>>& bufferViews);

	void SetValue(const std::string& name, RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset = 0, VkDeviceSize range = VK_WHOLE_SIZE);

	void SetValues(const std::string& name, const std::vector<RefCountPtr<RHIBuffer>>& buffers, std::vector<VkDeviceSize> offsets = std::vector<VkDeviceSize>(), std::vector<VkDeviceSize> ranges = std::vector<VkDeviceSize>());

	void SetValue(const std::string& name, void* value);

	void SetValue(const std::string& name, RefCountPtr<AccelerationStructure> topLevelAS);

	void SetValues(const std::string& name, const std::vector<RefCountPtr<AccelerationStructure>>& topLevelASs);

protected:

	void CheckSetIndex(uint32 setIndex);

	void SetSAMPLER(ShaderParam* shaderParam, RefCountPtr<RHISampler> sampler);

	void SetSAMPLERs(ShaderParam* shaderParam, const std::vector<RefCountPtr<RHISampler>>& samplers);

	void SetCOMBINED_IMAGE_SAMPLER(ShaderParam* shaderParam, RefCountPtr<RHITexture> texture, RefCountPtr<RHISampler> sampler, VkImageLayout layout);

	void SetCOMBINED_IMAGE_SAMPLERs(ShaderParam* shaderParam, const std::vector<RefCountPtr<RHITexture>>& textures, const std::vector<RefCountPtr<RHISampler>>& samplers, const std::vector<VkImageLayout>& layouts);

	void SetSAMPLED_IMAGE(ShaderParam* shaderParam, RefCountPtr<RHITexture> texture, VkImageLayout layout);

	void SetSAMPLED_IMAGEs(ShaderParam* shaderParam, const std::vector<RefCountPtr<RHITexture>>& textures, const std::vector<VkImageLayout>& layouts);

	void SetSTORAGE_IMAGE(ShaderParam* shaderParam, RefCountPtr<RHITexture> texture, VkImageLayout layout);

	void SetSTORAGE_IMAGEs(ShaderParam* shaderParam, const std::vector<RefCountPtr<RHITexture>>& textures, const std::vector<VkImageLayout>& layouts);

	void SetINPUT_ATTACHMENT(ShaderParam* shaderParam, RefCountPtr<RenderTarget> texture, VkImageLayout layout);

	void SetUNIFORM_TEXEL_BUFFER(ShaderParam* shaderParam, RefCountPtr<BufferView> bufferView);

	void SetUNIFORM_TEXEL_BUFFERs(ShaderParam* shaderParam, const std::vector<RefCountPtr<BufferView>>& bufferViews);

	void SetSTORAGE_TEXEL_BUFFER(ShaderParam* shaderParam, RefCountPtr<BufferView> bufferView);

	void SetSTORAGE_TEXEL_BUFFERs(ShaderParam* shaderParam, const std::vector<RefCountPtr<BufferView>>& bufferViews);

	void SetUNIFORM_BUFFER(ShaderParam* shaderParam, RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset, VkDeviceSize range);

	void SetUNIFORM_BUFFERs(ShaderParam* shaderParam, const std::vector<RefCountPtr<RHIBuffer>>& buffers, const std::vector<VkDeviceSize>& offsets, const std::vector<VkDeviceSize>& ranges);

	void SetSTORAGE_BUFFER(ShaderParam* shaderParam, RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset, VkDeviceSize range);

	void SetSTORAGE_BUFFERs(ShaderParam* shaderParam, const std::vector<RefCountPtr<RHIBuffer>>& buffers, const std::vector<VkDeviceSize>& offsets, const std::vector<VkDeviceSize>& ranges);

	void SetUNIFORM_BUFFER_DYNAMIC(ShaderParam* shaderParam, RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset, VkDeviceSize range);

	void SetUNIFORM_BUFFER_DYNAMICs(ShaderParam* shaderParam, const std::vector<RefCountPtr<RHIBuffer>>& buffers, const std::vector<VkDeviceSize>& offsets, const std::vector<VkDeviceSize>& ranges);

	void SetSTORAGE_BUFFER_DYNAMIC(ShaderParam* shaderParam, RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset, VkDeviceSize range);

	void SetSTORAGE_BUFFER_DYNAMICs(ShaderParam* shaderParam, const std::vector<RefCountPtr<RHIBuffer>>& buffers, const std::vector<VkDeviceSize>& offsets, const std::vector<VkDeviceSize>& ranges);

	void SetINLINE_UNIFORM_BLOCK_EXT(ShaderParam* shaderParam, void* value);

	void SetACCELERATION_STRUCTURE_KHR(ShaderParam* shaderParam, RefCountPtr<AccelerationStructure> topLevelAS);

	void SetACCELERATION_STRUCTURE_KHRs(ShaderParam* shaderParam, const std::vector<RefCountPtr<AccelerationStructure>>& topLevelASs);

	void SetValueForWrite(uint32 frameIndex, DescriptorNode& descriptorNode);

protected:

	std::unordered_map<std::string, ShaderParam*> mShaderParams;

	std::vector<ShaderPtr<RHIShader>> mShaderStages;

	std::map<uint32, uint32> mDescriptorCounts;

	std::map<uint32, bool> mDescriptorMultiples;

	std::vector<VkDescriptorSetLayout> mDescriptorSetLayouts;

	std::vector<VkPushConstantRange> mPushConstantRanges;

	VkPipelineLayout mPipelineLayout;

	VkDescriptorPoolInlineUniformBlockCreateInfoEXT mDescriptorPoolInlineUniformBlockCreateInfoEXT;

	VkDescriptorPool mDescriptorPool;

	std::map<uint32, std::map<uint32, DescriptorNode>> mDescriptorTable;

	std::vector<std::vector<VkDescriptorSet>> mDescriptorSetTable;

	uint32 mSetIndex;

	uint32 mOrderIndex;
};

template<typename ShaderType>
ShaderPtr<ShaderType>
RHIDescriptor::Builder::AddShader() noexcept
{
	ShaderPtr<ShaderType> shader(0);
	mImpl->shaderStages.push_back(shader);
	return shader;
}

template<typename ShaderType>
RHIDescriptor::Builder& RHIDescriptor::Builder::PushShader() noexcept
{
	ShaderPtr<ShaderType> shader(0);
	mImpl->shaderStages.push_back(shader);
	return *this;
}
